package ast

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"

	"golang.org/x/tools/go/ast/astutil"
)

type instrumenter struct {
	traceImport string
	tracePkg    string
	traceFunc   string
}

func New(traceImport, tracePkg, traceFunc string) *instrumenter {
	return &instrumenter{
		traceImport: traceImport,
		tracePkg:    tracePkg,
		traceFunc:   traceFunc,
	}
}

func (a instrumenter) Instrument(filename string) ([]byte, error) {
	fset := token.NewFileSet()
	curAst, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
	if err != nil {
		return nil, fmt.Errorf("error parsing %s: %s", filename, err)
	}

	if !hasFuncDecl(curAst) {
		return nil, nil
	}

	// add import declaration
	astutil.AddImport(fset, curAst, a.traceImport)

	a.addDeferTraceIntoFuncDecls(curAst)

	buf := &bytes.Buffer{}
	err = format.Node(buf, fset, curAst)
	if err != nil {
		return nil, fmt.Errorf("error formatting new code: %w", err)
	}

	return buf.Bytes(), nil
}

func hasFuncDecl(f *ast.File) bool {
	if len(f.Decls) == 0 {
		return false
	}

	for _, decl := range f.Decls {
		// 类型断言，检查 decl 是否是 *ast.FuncDecl 类型
		_, ok := decl.(*ast.FuncDecl)
		if ok {
			return true
		}
	}

	return false
}

func (a instrumenter) addDeferTraceIntoFuncDecls(f *ast.File) {
	for _, decl := range f.Decls { // 遍历所有声明语句
		// 类型断言，是否是 函数声明
		fd, ok := decl.(*ast.FuncDecl)
		if ok {
			// 如果是函数声明，则诸如跟踪代码
			a.addDeferStmt(fd)
		}
	}
}

// AST 解析​：

// stmt → *ast.DeferStmt（匹配）
// ds.Call.Fun → trace.Trace() → *ast.CallExpr（匹配）
// ce.Fun → trace.Trace → *ast.SelectorExpr（匹配）
// se.X → trace → *ast.Ident（匹配）
// x.Name == "trace" && se.Sel.Name == "Trace"（匹配）
// ​结果​：返回 false（已存在，不插入）。
func (a instrumenter) addDeferStmt(fd *ast.FuncDecl) (added bool) {
	stmts := fd.Body.List

	for _, stmt := range stmts {
		// 当前语句是否 defer 语句
		ds, ok := stmt.(*ast.DeferStmt)
		if !ok {
			continue
		}

		// 检查 defer 调用的函数是否是 CallExpr
		// Example:
		// defer trace.Trace()() // ds.Call.Fun 是 'trace.Trace()'， 属于 *ast.CallExpr
		// defer fmt.Println() 	// ds.Call.Fun 是 'fmt.Println'，属于 *ast.SelectorExpr （不匹配）
		ce, ok := ds.Call.Fun.(*ast.CallExpr)
		if !ok {
			continue
		}

		// 检查 CallExpr 的函数调用部分是否是 SelectorExpr （即 pkg.Func 形式）
		// Example:
		// trace.Trace() // ce.Fun 是 'trace.Trace()'，属于 *ast.SelectorExpr
		// someFunc() // ce.Fun 是 'someFunc'，属于 *ast.Ident （不匹配）
		se, ok := ce.Fun.(*ast.SelectorExpr)
		if !ok {
			continue
		}

		// 检查 SelectorExpr 的 X 部分是否是 Ident （这里是包名）
		// Example:
		// trace.Trace() // se.X 是 'trace'，属于 *ast.Ident
		// pkg.sub.Trace() // se.X 是 'pkg.sub'，属于 *ast.SelectorExpr（不匹配）
		x, ok := se.X.(*ast.Ident)
		if !ok {
			continue
		}

		// 检查包名和方法名是否匹配目标值，如 trace 和 Trace
		// Example:
		// defer trace.Trace() // x.Name = "trace", se.Sel.Name = "Trace"
		// defer log.Println() // x.Name = "log", se.Sel.Name = "Println"
		if (x.Name == a.tracePkg) && (se.Sel.Name == a.traceFunc) {
			// already exist, return
			return false
		}
	}

	// not found "defer tracePkg.traceFun()()"
	// add one
	ds := &ast.DeferStmt{
		Call: &ast.CallExpr{
			Fun: &ast.CallExpr{
				Fun: &ast.SelectorExpr{
					X: &ast.Ident{
						Name: a.tracePkg,
					},
					Sel: &ast.Ident{
						Name: a.traceFunc,
					},
				},
			},
		},
	}

	newList := make([]ast.Stmt, len(stmts)+1)
	copy(newList[1:], stmts)
	newList[0] = ds
	fd.Body.List = newList

	return true
}
